/*
 * KineticsSolver.cpp
 *
 *  Created on: 28 Jul 2011
 *      Author: Allan
 */

#include "KineticSolver.h"
#include "Utilities/Algorithms.h"
#include "Utilities/Assert.h"

KineticSolver::KineticSolver()
{}

KineticSolver::KineticSolver(const Multiphase& multiphase, const CanonicalReactionSystem& eReactions, const ReactionSystem& kReactions)
{
	Initialiaze(multiphase, eReactions, kReactions);
}

KineticSolver::~KineticSolver()
{}

void KineticSolver::Initialiaze(const Multiphase& multiphase, const CanonicalReactionSystem& eReactions, const ReactionSystem& kReactions)
{
	this->multiphase = multiphase;
	this->eReactions = eReactions;
	this->kReactions = kReactions;
	
	// The primary (jSpecies), equilibrium (eSpecies) and kinetic (kSpecies) species
	const vector<string> jSpecies = eReactions.GetPrimarySpecies();
	const vector<string> eSpecies = eReactions.GetSecondarySpecies();
	const vector<string> kSpecies = Difference(kReactions.GetSpecies(), jSpecies);
	
	// Set the indexes of the primary (jIndexes), equilibrium (eIndexes) and kinetic (kIndexes) species
	jIndexes = multiphase[jSpecies];
	eIndexes = multiphase[eSpecies];
	kIndexes = multiphase[kSpecies];
	
	// Set the number of primary (Nj), equilibrium (Ne) and kinetic (Nk) species
	Nj = jIndexes.size();
	Ne = eIndexes.size();
	Nk = kIndexes.size();
	
	// Initialise the auxiliar Eigen vectors with zero entries
	nj   = VectorXd::Zero(Nj);
	ne   = VectorXd::Zero(Ne);
	nk   = VectorXd::Zero(Nk);
	uj   = VectorXd::Zero(Nj); 
	ujnk = VectorXd::Zero(Nj + Nk);
	
	// Initialise the ODE solver
	odeSolver.Initialiaze(Nj + Nk);
	
	// Initialiaze the equilibrium solver
	eSolver.Initialiaze(multiphase, eReactions);
}

void KineticSolver::SetControlSolution(double epsAbs, double epsRel)
{
	odeSolver.SetControlSolution(epsAbs, epsRel);
}

bool KineticSolver::Iterate(double& t, double tEnd, double& dt, double T, double P, VectorXd& n)
{
	// Transfer the nj, ne, and nk part of the vector (n) to the vectors (nj), (ne), and (nk) respectively
	for(unsigned j = 0; j < Nj; ++j) nj[j] = n[jIndexes[j]];
	for(unsigned e = 0; e < Ne; ++e) ne[e] = n[eIndexes[e]];
	for(unsigned k = 0; k < Nk; ++k) nk[k] = n[kIndexes[k]];
	
	// The canonical stoichiometric matrix of the system of equilibrium reactions
	const MatrixXd& vej = eReactions.GetCanonicalStoichiometricMatrix();
	
	// Calculate the uj and nk part of the vector (ujnk) = [uj nk], where uj = nj + \sum_e vej * ne
	ujnk.segment(00, Nj) = nj + vej.transpose() * ne;
	ujnk.segment(Nj, Nk) = nk;
	
	// The kinetics problem, which contains the definition of the ODE function to advance the solution in time 
	KineticProblem problem(multiphase, eReactions, kReactions);
	
	// Set the necessary parameters of the kinetic problem object
	problem.SetParameters(T, P, n);
	
	// Use the ODE solver to calculate the vector (ujnk) at the new time (tn) = (t) + (dt)
	odeSolver.Iterate(problem, t, tEnd, dt, ujnk);
	
	// Transfer the uj and nk part of the vector (ujnk) = [uj nk] to the vectors (uj) and (nk) respectively
	uj = ujnk.segment(00, Nj);
	nk = ujnk.segment(Nj, Nk);
	
	// Update the nk part of the vector (n)
	for(unsigned k = 0; k < Nk; ++k) n[kIndexes[k]] = nk[k];
	
	// Solve the equilibrium problem with total composition of the primary species (uj)
	auto diagnostic = eSolver.Solve(uj, T, P, n);
	
	return diagnostic.converged;
}

bool KineticSolver::Solve(double& t, double tEnd, double& dt, double T, double P, VectorXd& n)
{
	while(t < tEnd)
		if(!Iterate(t, tEnd, dt, T, P, n))
			return false;
	
	return true;
}

KineticProblem::KineticProblem(const Multiphase& multiphase, const CanonicalReactionSystem& eReactions, const ReactionSystem& kReactions) :
multiphase(multiphase), eReactions(eReactions), kReactions(kReactions)
{
	// The primary (jSpecies) and kinetic (kSpecies) species
	const vector<string> jSpecies = eReactions.GetPrimarySpecies();
	const vector<string> kSpecies = Difference(kReactions.GetSpecies(), jSpecies);
	
	// Set the number of primary (Nj) and kinetic (Nk) species
	Nj = jSpecies.size();
	Nk = kSpecies.size();
	
	// Set the indexes of the kinetic species
	kIndexes = multiphase[kSpecies];
	
	// Assemble the stoichiometric matrices (ukj) and (ukk)
	ukj = kReactions.AssembleStoichiometricMatrix(jSpecies);
	ukk = kReactions.AssembleStoichiometricMatrix(kSpecies);
	
	// Initialiaze the total equilibrium solver (eSolver)
	eSolver.Initialiaze(multiphase, eReactions);
}

void KineticProblem::SetParameters(double T, double P, const VectorXd& n)
{
	this->T = T;
	this->P = P;
	this->n = n;
}

void KineticProblem::Function(double t, const VectorXd& ujnk, VectorXd& R)
{
	// Transfer the uj and nk part of the vector (ujnk) = [uj nk] to the vectors (uj) and (nk) respectively
	const VectorXd uj = ujnk.segment(00, Nj);
	const VectorXd nk = ujnk.segment(Nj, Nk);
	
	// Update the nk part of the vector (n)
	for(unsigned k = 0; k < Nk; ++k) n[kIndexes[k]] = nk[k];
	
	// Solve the equilibrium problem with total composition of the primary species (uj)
	auto diagnostic = eSolver.Solve(uj, T, P, n);
	
	// Calculate the activities of the species, now that the composition vector (n) has been determined
	const VectorXd a = multiphase.Activities(T, P, n);
	
	// Calculate the mass rates of the kinetic-controlled reactions with temperature (T), pressure (P), composition vector (n), and activity vector (a)
	const VectorXd rk = kReactions.KineticRates(T, P, n, a);
	
	// Calculate the right-hand side of the ordinary differential equations
	R.segment(00, Nj) = ukj.transpose() * rk;
	R.segment(Nj, Nk) = ukk.transpose() * rk;
}
